import torch
import torch.nn as nn
import pandas as pd

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class TitanicModel(nn.Module):
    def __init__(self):
        super(TitanicModel, self).__init__()
        self.linear1 = nn.Linear(26, 24)
        self.linear2 = nn.Linear(24, 16)
        self.linear3 = nn.Linear(16, 12)
        self.linear4 = nn.Linear(12, 8)
        self.linear5 = nn.Linear(8, 6)
        self.linear6 = nn.Linear(6, 4)
        self.linear7 = nn.Linear(4, 2)
        self.linear8 = nn.Linear(2, 1)
        self.act = nn.Sigmoid()

    def forward(self, x):
        x = self.act(self.linear1(x))
        x = self.act(self.linear2(x))
        x = self.act(self.linear3(x))
        x = self.act(self.linear4(x))
        x = self.act(self.linear5(x))
        x = self.act(self.linear6(x))
        x = self.act(self.linear7(x))
        x = self.linear8(x)
        return x


titanicModel = TitanicModel().to(device)
model_path = "../models/titanic_model.pth"
titanicModel.load_state_dict(torch.load(model_path))
titanicModel.eval()

test_path = "../data/Titanic/titanic_test3.csv"
df_test = pd.read_csv(test_path)
X_test = torch.tensor(df_test.iloc[:, 1:].values, dtype=torch.float32, device=device)

with torch.no_grad():
    logits = titanicModel(X_test)
    prob = torch.sigmoid(logits)
    pred_class = (prob > 0.5).to(torch.int32).squeeze(1).cpu().numpy()

submission = pd.DataFrame({
    "PassengerId": df_test["PassengerId"].astype(int).values,
    "Survived": pred_class.astype(int)
})
out_path = "../data/Titanic/titanic_submission.csv"
submission.to_csv(out_path, index=False)

print(f"Saved: {out_path}, shape={submission.shape}")
